/*
 * Copyright (c) 2002 - 2006 IBM Corporation.
 * All rights reserved. This program and the accompanying materials
 * are made available under the terms of the Eclipse Public License v1.0
 * which accompanies this distribution, and is available at
 * http://www.eclipse.org/legal/epl-v10.html
 *
 * Contributors:
 *     IBM Corporation - initial API and implementation
 */
package com.ibm.wala.util.collections;

import com.ibm.wala.util.intset.MutableIntSet;
import com.ibm.wala.util.intset.MutableSparseIntSet;
import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.WeakHashMap;

/** utilities for managing {@link Map}s */
public class MapUtil {
  /**
   * @param M a mapping from Object -&gt; Set
   * @return the Set corresponding to key in M; create one if needed
   * @throws IllegalArgumentException if M is null
   * @throws ClassCastException if the key is of an inappropriate type for this map (optional)
   * @throws NullPointerException if the specified key is null and this map does not permit null
   *     keys (optional)
   */
  public static <K, T> Set<T> findOrCreateSet(Map<K, Set<T>> M, K key) {
    return findOrCreateValue(M, key, () -> HashSetFactory.make(2));
  }

  /**
   * @throws ClassCastException if the key is of an inappropriate type for this map (optional)
   * @throws NullPointerException if the specified key is null and this map does not permit null
   *     keys (optional)
   */
  public static <K> MutableIntSet findOrCreateMutableIntSet(Map<K, MutableIntSet> M, K key) {
    return findOrCreateValue(M, key, MutableSparseIntSet::makeEmpty);
  }

  /**
   * @return the Collection corresponding to key in M; create one if needed
   * @throws ClassCastException if the key is of an inappropriate type for this map (optional)
   * @throws NullPointerException if the specified key is null and this map does not permit null
   *     keys (optional)
   */
  public static <K, T> Collection<T> findOrCreateCollection(Map<K, Collection<T>> M, K key) {
    return findOrCreateValue(M, key, () -> HashSetFactory.make(2));
  }

  /**
   * @return the Set corresponding to key in M; create one if needed
   * @throws IllegalArgumentException if M is null
   * @throws ClassCastException if the key is of an inappropriate type for this map (optional)
   * @throws NullPointerException if the specified key is null and this map does not permit null
   *     keys (optional)
   */
  public static <K, T> List<T> findOrCreateList(Map<K, List<T>> M, K key) {
    return findOrCreateValue(M, key, ArrayList::new);
  }

  /**
   * @param M a mapping from Object -&gt; Map
   * @return the Map corresponding to key in M; create one if needed
   * @throws IllegalArgumentException if M is null
   * @throws ClassCastException if the key is of an inappropriate type for this map (optional)
   * @throws NullPointerException if the specified key is null and this map does not permit null
   *     keys (optional)
   */
  public static <K, K2, V> Map<K2, V> findOrCreateMap(Map<K, Map<K2, V>> M, K key) {
    return findOrCreateValue(M, key, () -> HashMapFactory.make(2));
  }

  /**
   * @throws ClassCastException if the key is of an inappropriate type for this map (optional)
   * @throws NullPointerException if the specified key is null and this map does not permit null
   *     keys (optional)
   */
  public static <K, V> V findOrCreateValue(Map<K, V> M, K key, Factory<V> factory) {
    if (M == null) {
      throw new IllegalArgumentException("M is null");
    }
    return M.computeIfAbsent(key, missing -> factory.make());
  }

  /**
   * @param M a mapping from Object -&gt; WeakHashMap
   * @return the WeakHashMap corresponding to key in M; create one if needed
   * @throws IllegalArgumentException if M is null
   * @throws ClassCastException if the key is of an inappropriate type for this map (optional)
   * @throws NullPointerException if the specified key is null and this map does not permit null
   *     keys (optional)
   */
  public static <K, V> WeakHashMap<K, V> findOrCreateWeakHashMap(
      Map<Object, WeakHashMap<K, V>> M, Object key) {
    return findOrCreateValue(M, key, () -> new WeakHashMap<>(2));
  }

  /**
   * @param m a map from key -&gt; {@link Set}&lt;value&gt;
   * @return inverted map, value -&gt; {@link Set}&lt;key&gt;
   * @throws IllegalArgumentException if m is null
   */
  public static <K, V> Map<V, Set<K>> inverseMap(Map<K, Set<V>> m) {
    if (m == null) {
      throw new IllegalArgumentException("m is null");
    }
    Map<V, Set<K>> result = HashMapFactory.make(m.size());
    m.forEach((key, values) -> values.forEach(value -> findOrCreateSet(result, value).add(key)));
    return result;
  }

  /**
   * invert an input map that is one-to-one (i.e., it does not map two different keys to the same
   * value)
   *
   * @throws IllegalArgumentException if m is null
   * @throws IllegalArgumentException if m is not one-to-one
   */
  public static <K, V> Map<V, K> invertOneToOneMap(Map<K, V> m) {
    if (m == null) {
      throw new IllegalArgumentException("m is null");
    }
    Map<V, K> result = HashMapFactory.make(m.size());
    m.forEach(
        (key, value) ->
            result.merge(
                value,
                key,
                (oldValue, newValue) -> {
                  throw new IllegalArgumentException("input map not one-to-one");
                }));
    return result;
  }

  public static <K, V> Map<Set<K>, V> groupKeysByValue(Map<K, V> m) {
    if (m == null) {
      throw new IllegalArgumentException("m is null");
    }
    Map<V, Set<K>> valueToKeys = HashMapFactory.make();
    m.forEach((key, value) -> findOrCreateSet(valueToKeys, value).add(key));
    return invertOneToOneMap(valueToKeys);
  }
}
